#!/usr/bin/python
import sys
import os
import string
import math
from openeye.oechem import *

if len(sys.argv) <> 6:
	print 'pharmacophore_search_on_docking_hits.py <sybyl hits-file with molecule names> <pdb-file with pharmacophor-sites> <site radius> <number of sites to match> <output-file>'
	#print 'don\'t forget to execute within .mdb-folder!'
	sys.exit()

#pharmacophor sites (lig)
l_x_don = []
l_y_don = []
l_z_don = []
l_x_acc = []
l_y_acc = []
l_z_acc = []
radius = float(sys.argv[3]) # 0.8-1.2
pharm = open(sys.argv[2], 'r')
for i in pharm.readlines():
	if i[17:20] == 'DON':
		l_x_don.append(float(i[31:38].strip()))
		l_y_don.append(float(i[39:46].strip()))
		l_z_don.append(float(i[47:54].strip()))
	if i[17:20] == 'ACC':
		l_x_acc.append(float(i[31:38].strip()))
		l_y_acc.append(float(i[39:46].strip()))
		l_z_acc.append(float(i[47:54].strip()))
pharm.close()
i_acc = len(l_x_acc)
i_don = len(l_x_don)
print i_acc, 'acceptors found,', i_don, 'donors found.'

ifs = oemolistream()
mol = OEGraphMol()

files = []
file = open(sys.argv[1], 'r')
for i in file.readlines():
	if 'name="' in i:
		files.append('pharm.mdb/' + i.split('name="')[1].split('";')[0] + '.mol2')

#files = os.listdir('./')
#alle mol2 files in directory, remove all files which are not mol2
#files.sort()
#tmp_list = []
#for i in files:
#  if i[-4:] == 'mol2':
#    tmp_list.append(i)
#files = tmp_list

ergebnis = open(sys.argv[5],'w')


donor_list = ['S.3','O.3','N.am','N.ar','N.pl3','N.4','N.3','N.2','C.2','C.ar']
#donor_liste = ['O.3']
acceptor_list =['S.3','S.2','O.2','O.co2','N.1','O.3','N.ar','N.3','N.2','Cl','Br','I']
#acceptor_list =['O.2','S.2','O.co2','N.1']

#fuer alle files

print "scanning files"

for i in files:
   print i
   ifs.open(i)
   OEReadMolecule(ifs, mol)
   OEAssignAromaticFlags(mol,OEAroModelTripos)
   OETriposAtomNames(mol)
   OETriposAtomTypeNames(mol)
   OETriposBondTypeNames(mol)
   OETriposAtomNames(mol)
   
   #print mol.GetTitle()
   
   donor_found = 0
   acceptor_found = 0
   
   for atom in mol.GetAtoms():
   
     atom_type = atom.GetType()
     if atom_type in donor_list and donor_found < i_don:
       #check if is is connected to an H
       donor = False
       donor_h = False
       donor_n = False
       for nbor in atom.GetAtoms():
         if atom_type <> 'C.ar' and atom_type <> 'C.2':
           if  nbor.GetType() == 'H':
             donor = True
         else:
           if  nbor.GetType() == 'H':
             donor_h = True
           if nbor.GetType() == 'N.ar' or nbor.GetType() == 'N.2' or nbor.GetType() == 'N.pl3':
             donor_n = True
         if donor_h and donor_n:
             donor = True
       if donor == True:
           #print atom_type
	   #check if coordinates are okay
           x = mol.GetCoords(atom)[0]
	   y = mol.GetCoords(atom)[1]
	   z = mol.GetCoords(atom)[2]
           i_counter = 0
           for site in l_x_don:
              distance = math.sqrt((x-site)**2+(y-l_y_don[i_counter])**2+(z-l_z_don[i_counter])**2)
              i_counter += 1
              if distance <= radius:
	         donor_found += 1 
              #print distance
              #print donor_found
     elif atom_type in acceptor_list and acceptor_found < i_acc:
           #print atom_type
	   #check if coordinates are okay
	   x = mol.GetCoords(atom)[0]
	   y = mol.GetCoords(atom)[1]
	   z = mol.GetCoords(atom)[2]
           i_counter = 0
           for site in l_x_acc:
              distance = math.sqrt((x-site)**2+(y-l_y_acc[i_counter])**2+(z-l_z_acc[i_counter])**2)
              i_counter += 1
              if distance <= radius:
	         acceptor_found += 1 
              #print distance
              #print donor_found

   if donor_found + acceptor_found >= int(sys.argv[4]):   
      ergebnis.write(mol.GetTitle() + '\tDON: ' + str(donor_found) + '\tACC: ' + str(acceptor_found) + '\n')
   
ergebnis.close()      
  
